package com.jscloud.spark.broadcast

import org.apache.spark.rdd.RDD
import org.apache.spark.util.AccumulatorV2
import org.apache.spark.{SparkConf, SparkContext}

object Accumulator02ForUdf {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf().setAppName("Accumulator02ForUdf").setMaster("local[*]")
    val sc = new SparkContext(conf)
    sc.setLogLevel("WARN")

    val rdd1: RDD[Int] = sc.parallelize(Array(1, 2, 3, 4, 5), 2)
    //自定义累加器
    val acc: MyAccumulator = new MyAccumulator

    //注册累加器
    sc.register(acc)

    rdd1.foreach(x => {
      acc.add(1)
      println("rdd:" + acc.value)
    })
    println("------------")
    println("main" + acc.value)

    sc.stop()
  }
}

class MyAccumulator extends AccumulatorV2[Int, Int] {
  var sum: Int = 0

  override def isZero: Boolean = sum == 0

  override def copy(): AccumulatorV2[Int, Int] = {
    val accumulator: MyAccumulator = new MyAccumulator
    accumulator.sum = sum
    accumulator
  }

  override def reset(): Unit = {
    sum = 0
  }

  override def add(v: Int): Unit = {
    sum += v
  }

  override def merge(other: AccumulatorV2[Int, Int]): Unit = {
    other match {
      case o: MyAccumulator => this.sum += o.sum
      case _ =>
    }
  }

  override def value: Int = {
    this.sum
  }
}
